{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "## 把数据处理成模型需要的文件格式"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import struct\n",
    "import collections\n",
    "from tensorflow.core.example import example_pb2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 经过分词处理后的训练数据与测试数据文件\n",
    "TRAIN_FILE = \"./data/weibo_news/train_art_summ_prep.txt\"\n",
    "VAL_FILE = \"./data/weibo_news/val_art_summ_prep.txt\"\n",
    "\n",
    "# 文本起始与结束标志\n",
    "SENTENCE_START = '<s>'\n",
    "SENTENCE_END = '</s>'\n",
    "\n",
    "VOCAB_SIZE = 50_000  # 词汇表大小\n",
    "CHUNK_SIZE = 1000    # 每个分块example的数量，用于分块的数据\n",
    "\n",
    "# tf模型数据文件存放目录\n",
    "FINISHED_FILE_DIR = './data/weibo_news/finished_files'\n",
    "CHUNKS_DIR = os.path.join(FINISHED_FILE_DIR, 'chunked')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def chunk_file(finished_files_dir, chunks_dir, name, chunk_size):\n",
    "    \"\"\"构建二进制文件\"\"\"\n",
    "    in_file = os.path.join(finished_files_dir, '%s.bin' % name)\n",
    "    print(in_file)\n",
    "    reader = open(in_file, \"rb\")\n",
    "    chunk = 0\n",
    "    finished = False\n",
    "    while not finished:\n",
    "        chunk_fname = os.path.join(chunks_dir, '%s_%03d.bin' % (name, chunk))  # 新的分块\n",
    "        with open(chunk_fname, 'wb') as writer:\n",
    "            for _ in range(chunk_size):\n",
    "                len_bytes = reader.read(8)\n",
    "                if not len_bytes:\n",
    "                    finished = True\n",
    "                    break\n",
    "                str_len = struct.unpack('q', len_bytes)[0]\n",
    "                example_str = struct.unpack('%ds' % str_len, reader.read(str_len))[0]\n",
    "                writer.write(struct.pack('q', str_len))\n",
    "                writer.write(struct.pack('%ds' % str_len, example_str))\n",
    "            chunk += 1\n",
    "\n",
    "\n",
    "def chunk_all():\n",
    "    # 创建一个文件夹来保存分块\n",
    "    if not os.path.isdir(CHUNKS_DIR):\n",
    "        os.mkdir(CHUNKS_DIR)\n",
    "    # 将数据分块\n",
    "    for name in ['train', 'val']:\n",
    "        print(\"Splitting %s data into chunks...\" % name)\n",
    "        chunk_file(FINISHED_FILE_DIR, CHUNKS_DIR, name, CHUNK_SIZE)\n",
    "    print(\"Saved chunked data in %s\" % CHUNKS_DIR)\n",
    "\n",
    "\n",
    "def read_text_file(text_file):\n",
    "    \"\"\"从预处理好的文件中加载数据\"\"\"\n",
    "    lines = []\n",
    "    with open(text_file, \"r\", encoding='utf-8') as f:\n",
    "        for line in f:\n",
    "            lines.append(line.strip())\n",
    "    return lines\n",
    "\n",
    "\n",
    "def write_to_bin(input_file, out_file, makevocab=False):\n",
    "    \"\"\"生成模型需要的文件\"\"\"\n",
    "    if makevocab:\n",
    "        vocab_counter = collections.Counter()\n",
    "\n",
    "    with open(out_file, 'wb') as writer:\n",
    "        # 读取输入的文本文件，使偶数行成为article，奇数行成为abstract（行号从0开始）\n",
    "        lines = read_text_file(input_file)\n",
    "        for i, new_line in enumerate(lines):\n",
    "            if i % 2 == 0:\n",
    "                article = lines[i]\n",
    "            if i % 2 != 0:\n",
    "                abstract = \"%s %s %s\" % (SENTENCE_START, lines[i], SENTENCE_END)\n",
    "\n",
    "                # 写入tf.Example\n",
    "                tf_example = example_pb2.Example()\n",
    "                tf_example.features.feature['article'].bytes_list.value.extend([bytes(article, encoding='utf-8')])\n",
    "                tf_example.features.feature['abstract'].bytes_list.value.extend([bytes(abstract, encoding='utf-8')])\n",
    "                tf_example_str = tf_example.SerializeToString()\n",
    "                str_len = len(tf_example_str)\n",
    "                writer.write(struct.pack('q', str_len))\n",
    "                writer.write(struct.pack('%ds' % str_len, tf_example_str))\n",
    "\n",
    "                # 如果可以，将词典写入文件\n",
    "                if makevocab:\n",
    "                    art_tokens = article.split(' ')\n",
    "                    abs_tokens = abstract.split(' ')\n",
    "                    abs_tokens = [t for t in abs_tokens if\n",
    "                                  t not in [SENTENCE_START, SENTENCE_END]]  # 从词典中删除这些符号\n",
    "                    tokens = art_tokens + abs_tokens\n",
    "                    tokens = [t.strip() for t in tokens]     # 去掉句子开头结尾的空字符\n",
    "                    tokens = [t for t in tokens if t != \"\"]  # 删除空行\n",
    "                    vocab_counter.update(tokens)\n",
    "\n",
    "    print(\"Finished writing file %s\\n\" % out_file)\n",
    "\n",
    "    # 将词典写入文件\n",
    "    if makevocab:\n",
    "        print(\"Writing vocab file...\")\n",
    "        with open(os.path.join(FINISHED_FILE_DIR, \"vocab\"), 'w', encoding='utf-8') as writer:\n",
    "            for word, count in vocab_counter.most_common(VOCAB_SIZE):\n",
    "                writer.write(word + ' ' + str(count) + '\\n')\n",
    "        print(\"Finished writing vocab file\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Finished writing file ./data/weibo_news/finished_files/val.bin\n",
      "\n",
      "Finished writing file ./data/weibo_news/finished_files/train.bin\n",
      "\n",
      "Writing vocab file...\n",
      "Finished writing vocab file\n"
     ]
    },
    {
     "ename": "TypeError",
     "evalue": "chunk_all() missing 1 required positional argument: 'direction'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-4-97fb98050243>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      4\u001b[0m \u001b[0mwrite_to_bin\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mTRAIN_FILE\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mos\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpath\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mjoin\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mFINISHED_FILE_DIR\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m\"train.bin\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmakevocab\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 6\u001b[0;31m \u001b[0mchunk_all\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;31mTypeError\u001b[0m: chunk_all() missing 1 required positional argument: 'direction'"
     ]
    }
   ],
   "source": [
    "if not os.path.exists(FINISHED_FILE_DIR):\n",
    "    os.makedirs(FINISHED_FILE_DIR)\n",
    "write_to_bin(VAL_FILE, os.path.join(FINISHED_FILE_DIR, \"val.bin\"))\n",
    "write_to_bin(TRAIN_FILE, os.path.join(FINISHED_FILE_DIR, \"train.bin\"), makevocab=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Splitting train data into chunks...\n",
      "./data/weibo_news/finished_files/train.bin\n",
      "Splitting val data into chunks...\n",
      "./data/weibo_news/finished_files/val.bin\n",
      "Saved chunked data in ./data/weibo_news/finished_files/chunked\n"
     ]
    }
   ],
   "source": [
    "chunk_all()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "## struct tf.exam"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
